from fastapi import APIRouter, Request, File, HTTPException, Depends, Form
from module_admin.service.key_service import KeyService
from pydantic import BaseModel
from utils.response_util import ResponseUtil
from utils.ocr_llm_util import ImageOptimizer, OcrLlmConfig
from ocr_core.schema import OcrTask, OcrTaskResult, OcrLlmTaskResult, OcrLlmModel
from ocr_core.ocr_producer import OcrProducer
from ocr_core.llm_producer import LlmProducer
import base64
import asyncio
import uuid
import json
from config.mongodb import TaskMDB

OcrController = APIRouter(
    prefix="/ocr",
    dependencies=[
        Depends(KeyService.key_is_exist),
        Depends(OcrLlmConfig.scenario_check),
    ],
)


# -------------------- JSON 字符串类型 --------------------
class JsonString(str):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v, field):
        if v == "" or v is None:
            return v
        try:
            json.loads(v)
        except Exception:
            raise ValueError(f"必须是合法的 JSON 字符串")
        return str(v)


# -------------------- 请求体模型 --------------------
class ImageUrlRequest(BaseModel):
    image_url: str
    prompt: JsonString = ""
    extra: JsonString = ""


class Base64ImageRequest(BaseModel):
    image_base64: str
    prompt: JsonString = ""
    extra: JsonString = ""


# -------------------- 工具函数 --------------------
async def save_base64_image(image_base64: str) -> str:
    """保存 base64 图片为本地文件"""
    if "base64," in image_base64:
        _, base64_data = image_base64.split("base64,", 1)
    else:
        base64_data = image_base64
    try:
        image_data = base64.b64decode(base64_data)
    except Exception:
        raise HTTPException(status_code=400, detail="无效的base64图像数据")

    file_path = ImageOptimizer.image_path(upload=True)
    with open(file_path, "wb") as f:
        f.write(image_data)
    return file_path


async def create_task(
    request: Request, image_path: str, llm: bool, scenario: str, prompt: str, extra: str
) -> str:
    """统一创建 OCR / OCR+LLM 任务"""
    task = OcrTask(
        task_id=str(uuid.uuid4()).replace("-", ""),
        image_path=image_path,
        optimized_path="",
        llm=llm,
        scenario=scenario,
        user_prompt=prompt,
        extra=extra,
    )
    producer = OcrProducer(request.app.state.redis)
    return await producer.create_ocr_task(task)


async def wait_for_completion(
    request: Request,
    task_id: str,
    llm: bool = False,
    max_retries: int = 60,
    interval: float = 1.0,
) -> dict:
    """
    轮询任务直到完成
    - llm = False → 纯 OCR
    - llm = True  → OCR + LLM
    """
    ocr_producer = OcrProducer(request.app.state.redis)
    llm_producer = LlmProducer(request.app.state.redis) if llm else None

    ocr_task, llm_task, task_status = None, None, "pending"

    for _ in range(max_retries):
        await asyncio.sleep(interval)

        ocr_task = await ocr_producer.get_ocr_task_result(task_id)
        task_status = ocr_task.status

        if llm:
            llm_task = await llm_producer.get_llm_task_result(task_id)
            # 合并状态
            if ocr_task.status == "completed":
                task_status = (
                    llm_task.status
                    if llm_task.status in ["completed", "failed"]
                    else "processing"
                )

        if task_status in ["completed", "failed"]:
            break

    result = ocr_task.model_dump()
    if llm and llm_task:
        result.update(llm_task.model_dump())
    result["status"] = task_status
    return result


# -------------------- 通用接口生成器 --------------------
def register_task_routes(llm: bool, sync: bool):
    """
    自动注册接口，避免重复代码
    - llm: 是否带 LLM
    - sync: 是否同步（轮询完成）
    """

    suffix = "llm/" if llm else ""
    mode = "" if sync else "-task"
    wait_func = (
        (lambda req, tid: wait_for_completion(req, tid, llm=llm)) if sync else None
    )

    @OcrController.post(f"/{suffix}url{mode}", response_model=OcrLlmModel)
    async def url_endpoint(request: Request, body: ImageUrlRequest):
        scenario = request.query_params.get("scenario", "material_label")
        task_id = await create_task(
            request,
            body.image_url,
            llm=llm,
            scenario=scenario,
            prompt=body.prompt,
            extra=body.extra,
        )
        if sync:
            return ResponseUtil.success(data=await wait_func(request, task_id))
        return ResponseUtil.success(data={"task_id": task_id})

    @OcrController.post(f"/{suffix}base64{mode}", response_model=OcrLlmModel)
    async def base64_endpoint(request: Request, body: Base64ImageRequest):
        scenario = request.query_params.get("scenario", "material_label")
        file_path = await save_base64_image(body.image_base64)
        task_id = await create_task(
            request,
            file_path,
            llm=llm,
            scenario=scenario,
            prompt=body.prompt,
            extra=body.extra,
        )
        if sync:
            return ResponseUtil.success(data=await wait_func(request, task_id))
        return ResponseUtil.success(data={"task_id": task_id})

    @OcrController.post(f"/{suffix}binary{mode}", response_model=OcrLlmModel)
    async def binary_endpoint(
        request: Request,
        file: bytes = File(...),
        prompt: JsonString = Form(""),
        extra: JsonString = Form(""),
    ):
        scenario = request.query_params.get("scenario", "material_label")
        file_path = ImageOptimizer.image_path(upload=True)
        with open(file_path, "wb") as f:
            f.write(file)
        task_id = await create_task(
            request, file_path, llm=llm, scenario=scenario, prompt=prompt, extra=extra
        )
        if sync:
            return ResponseUtil.success(data=await wait_func(request, task_id))
        return ResponseUtil.success(data={"task_id": task_id})


# -------------------- 注册所有接口 --------------------
# OCR 异步
register_task_routes(llm=False, sync=False)
# OCR 同步
register_task_routes(llm=False, sync=True)
# OCR+LLM 异步
register_task_routes(llm=True, sync=False)
# OCR+LLM 同步
register_task_routes(llm=True, sync=True)


# -------------------- 查询接口保留 --------------------
@OcrController.get("/task/{task_id}", response_model=OcrTaskResult)
async def get_ocr_task_result(request: Request, task_id: str):
    """获取OCR任务结果"""
    ocr_producer = OcrProducer(request.app.state.redis)
    ocr_task = await ocr_producer.get_ocr_task_result(task_id)
    return ResponseUtil.success(data=ocr_task)


@OcrController.get("/llm/task/{task_id}", response_model=OcrLlmTaskResult)
async def get_ocr_llm_task_result(request: Request, task_id: str):
    """获取OCR+LLM任务结果"""
    ocr_producer = OcrProducer(request.app.state.redis)
    llm_producer = LlmProducer(request.app.state.redis)
    ocr_task = await ocr_producer.get_ocr_task_result(task_id)
    llm_task = await llm_producer.get_llm_task_result(task_id)

    merged_status = ocr_task.status
    if ocr_task.status == "completed":
        merged_status = (
            llm_task.status
            if llm_task.status in ["completed", "failed"]
            else "processing"
        )

    return ResponseUtil.success(
        data={**ocr_task.model_dump(), **llm_task.model_dump(), "status": merged_status}
    )
